# 用于输出rating_title.csv
from mongo import (
  collection_movie_info,
  get_all_data
)
from utils import (
  get_data_from_json,
  get_data_from_csv,
  log_message,
  store_data_to_csv
)

def create_rating_title_csv():
  result = []
  # 1.拿到tmdb数据
  movie_data = get_all_data(collection_movie_info)
  # 2.拿到rating数据
  log_message("正在获取rating.csv数据")
  rating_data = get_data_from_csv(path='../data/origin-data/', file_name='ratings.csv')
  # 3.拿到rating与tmdb关联的数据
  link_data = get_data_from_json(filename='links.json')

  # 4.创建tmdbId到title的映射
  tmdb_to_title = {}
  for movie in movie_data:
    if "title" in movie:
      tmdb_to_title[movie["id"]] = movie["title"]

  # 5.创建movieId到tmdbId的映射
  movie_to_tmdb = {link["movieId"]: int(link["tmdbId"]) for link in link_data}
  log_message("开始构造数据")

  count = 0
  totle = len(rating_data)
  # 6.构造需要存储的数据
  for rating in rating_data:
    movieId = rating["movieId"]
    if movieId in movie_to_tmdb:
      tmdbId = movie_to_tmdb[movieId]
      if tmdbId in tmdb_to_title:
        title = tmdb_to_title[tmdbId]
        result.append({
          "userId": rating["userId"],
          "movieId": movieId,
          "rating": rating["rating"],
          "timestamp": rating["timestamp"],
          "tmdbId": tmdbId,
          "title": title
        })
    count += 1
    if count % 10000 == 0:
      log_message(f"已处理{round(count / totle * 100, 2)}%")
      store_data_to_csv(result, 'ratings_title.csv')
      result = []
  if len(result) > 0:
    store_data_to_csv(result, 'ratings_title.csv')
    result = []

create_rating_title_csv()


